package intelligentCustomer.memory;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 基于向量数据库的记忆存储库
 * 使用嵌入向量存储和检索用户记忆
 */
public class VectorMemoryRepository implements MemoryRepository {
    /** 嵌入模型 */
    private final EmbeddingModel embeddingModel;
    
    /** 嵌入存储 */
    private final EmbeddingStore<TextSegment> embeddingStore;
    
    /** 文档分割器 */
    private final DocumentSplitter documentSplitter;
    
    /** 用户档案缓存 */
    private final Map<String, UserProfile> profileCache = new ConcurrentHashMap<>();
    
    /** 日期时间格式化器 */
    private static final DateTimeFormatter FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
    
    /**
     * 构造函数
     */
    public VectorMemoryRepository() {
        // 初始化嵌入模型
        this.embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey("sk-bde931a661f74ba4ae96ad4050299a46")
                .baseUrl("https://api.deepseek.com/v1")
                .modelName("deepseek-embedding")
                .build();
        
        // 初始化嵌入存储
        this.embeddingStore = new InMemoryEmbeddingStore<>();
        
        // 初始化文档分割器
        this.documentSplitter = DocumentSplitters.recursive(100, 0);
    }
    
    /**
     * 加载用户档案
     * 
     * @param userId 用户ID
     * @return 用户档案对象
     */
    @Override
    public UserProfile loadUserProfile(String userId) {
        // 检查缓存
        if (profileCache.containsKey(userId)) {
            return profileCache.get(userId);
        }
        
        // 创建新的用户档案
        UserProfile profile = new UserProfile(userId);
        
        // 从向量数据库检索用户相关记忆
        List<EmbeddingMatch<TextSegment>> relevantMemories = retrieveRelevantMemories(userId, 10);
        
        // 解析记忆并更新用户档案
        for (EmbeddingMatch<TextSegment> match : relevantMemories) {
            String memoryText = match.embedded().text();
            parseAndUpdateProfile(profile, memoryText);
        }
        
        // 缓存用户档案
        profileCache.put(userId, profile);
        
        return profile;
    }
    
    /**
     * 保存用户档案
     * 
     * @param profile 用户档案
     * @return 是否保存成功
     */
    @Override
    public boolean saveUserProfile(UserProfile profile) {
        try {
            // 更新缓存
            profileCache.put(profile.getUserId(), profile);
            
            // 将用户档案转换为记忆文本
            String memoryText = convertProfileToMemoryText(profile);
            
            // 创建文档
            Document document = Document.from(memoryText);
            
            // 分割文档
            List<TextSegment> segments = documentSplitter.split(document);
            
            // 为每个段落生成嵌入并存储
            for (TextSegment segment : segments) {
                Embedding embedding = embeddingModel.embed(segment).content();
                embeddingStore.add(embedding, segment);
            }
            
            return true;
        } catch (Exception e) {
            System.err.println("保存用户档案到向量数据库失败: " + e.getMessage());
            return false;
        }
    }
    
    /**
     * 保存对话记忆
     * 
     * @param userId 用户ID
     * @param role 角色
     * @param content 内容
     * @param timestamp 时间戳
     * @return 是否保存成功
     */
    public boolean saveConversationMemory(String userId, String role, String content, LocalDateTime timestamp) {
        try {
            // 构建记忆文本
            String memoryText = String.format("USER_ID: %s\nTYPE: conversation\nROLE: %s\nCONTENT: %s\nTIMESTAMP: %s",
                    userId, role, content, timestamp.format(FORMATTER));
            
            // 创建文档
            Document document = Document.from(memoryText);
            
            // 分割文档
            List<TextSegment> segments = documentSplitter.split(document);
            
            // 为每个段落生成嵌入并存储
            for (TextSegment segment : segments) {
                Embedding embedding = embeddingModel.embed(segment).content();
                embeddingStore.add(embedding, segment);
            }
            
            return true;
        } catch (Exception e) {
            System.err.println("保存对话记忆到向量数据库失败: " + e.getMessage());
            return false;
        }
    }
    
    /**
     * 检索相关记忆
     * 
     * @param query 查询文本
     * @param maxResults 最大结果数
     * @return 相关记忆列表
     */
    public List<EmbeddingMatch<TextSegment>> retrieveRelevantMemories(String query, int maxResults) {
        try {
            // 生成查询嵌入
            Embedding queryEmbedding = embeddingModel.embed(query).content();
            
            // 在向量数据库中搜索相似嵌入
            return embeddingStore.findRelevant(queryEmbedding, maxResults);
        } catch (Exception e) {
            System.err.println("从向量数据库检索记忆失败: " + e.getMessage());
            return new ArrayList<>();
        }
    }
    
    /**
     * 将用户档案转换为记忆文本
     * 
     * @param profile 用户档案
     * @return 记忆文本
     */
    private String convertProfileToMemoryText(UserProfile profile) {
        StringBuilder sb = new StringBuilder();
        
        sb.append("USER_ID: ").append(profile.getUserId()).append("\n");
        sb.append("TYPE: profile\n");
        
        // 添加产品偏好
        sb.append("PRODUCT_PREFERENCES: ");
        List<String> preferences = profile.getProductPreferences();
        if (!preferences.isEmpty()) {
            for (String pref : preferences) {
                sb.append(pref).append(", ");
            }
        } else {
            sb.append("none");
        }
        sb.append("\n");
        
        // 添加价格敏感度
        sb.append("PRICE_SENSITIVITY: ").append(profile.getPriceSensitivity()).append("\n");
        
        // 添加查询次数
        sb.append("QUERY_COUNT: ").append(profile.getQueryCount()).append("\n");
        
        // 添加首次交互时间
        sb.append("FIRST_INTERACTION: ").append(profile.getFirstInteractionTime().format(FORMATTER)).append("\n");
        
        // 添加最近交互时间
        sb.append("LAST_INTERACTION: ").append(profile.getLastInteractionTime().format(FORMATTER)).append("\n");
        
        return sb.toString();
    }
    
    /**
     * 解析记忆文本并更新用户档案
     * 
     * @param profile 用户档案
     * @param memoryText 记忆文本
     */
    private void parseAndUpdateProfile(UserProfile profile, String memoryText) {
        // 解析记忆文本中的各种属性
        if (memoryText.contains("PRODUCT_PREFERENCES:")) {
            String prefsText = memoryText.substring(memoryText.indexOf("PRODUCT_PREFERENCES:") + 20);
            prefsText = prefsText.substring(0, prefsText.indexOf("\n"));
            String[] prefs = prefsText.split(", ");
            for (String pref : prefs) {
                if (!pref.equals("none")) {
                    profile.addProductPreference(pref.trim());
                }
            }
        }
        
        if (memoryText.contains("PRICE_SENSITIVITY:")) {
            String sensitivity = memoryText.substring(memoryText.indexOf("PRICE_SENSITIVITY:") + 18);
            sensitivity = sensitivity.substring(0, sensitivity.indexOf("\n")).trim();
            profile.setPriceSensitivity(sensitivity);
        }
    }

    /**
     * 保存用户记忆
     * 
     * @param userId 用户ID
     * @param content 内容
     * @param timestamp 时间戳
     */
    public void saveMemory(String userId, String content, LocalDateTime timestamp) {
        // 实现保存记忆的逻辑
        // 可以调用已有的saveConversationMemory方法
        saveConversationMemory(userId, "用户", content, timestamp);
    }
}